package interceptor

import (
	"context"
	"github.com/zeromicro/go-zero/core/logx"
	"google.golang.org/grpc"
	"mall/common/excption"
	"mall/common/limit"
)

type RateLimitInterceptor struct {
	limit        *limit.PeriodLimit
	rateLimitKey string
}

func NewRateLimitInterceptor(periodLimit *limit.PeriodLimit, rateLimitKey string) *RateLimitInterceptor {
	return &RateLimitInterceptor{limit: periodLimit, rateLimitKey: rateLimitKey}
}

// RateLimitInterceptor 统一限流配置
func (r *RateLimitInterceptor) RateLimitInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
	handler grpc.UnaryHandler) (resp interface{}, err error) {

	takeCode, err := r.limit.TakeCtx(ctx, r.rateLimitKey)
	if err != nil {
		logx.Error(err)
		return true, err
	}
	limitErr := doRateLimit(takeCode, r.rateLimitKey)
	if limitErr != nil {
		return nil, limitErr
	}
	return handler(ctx, req)
}

func doRateLimit(takeCode int, rateLimitKey string) error {
	switch takeCode {
	case limit.OverQuota:
		logx.Errorf("OverQuota key: %v", rateLimitKey)
		return excption.NewErrCode(excption.RequestLimitError)
	case limit.Allowed:
		fallthrough
	case limit.HitQuota:
		logx.Infof("AllowedQuota key: %v", rateLimitKey)
		return nil
	default:
		logx.Errorf("DefaultQuota key: %v", rateLimitKey)
		// unknown response, we just let the sms go
		return excption.NewErrCode(excption.RequestLimitError)
	}
}
